import numpy as np
from PIL import Image
from torch.utils.data.dataset import Dataset
from torchvision import transforms
import torch
import os
# In[] 用于读取图片
class CustomDataset(Dataset):
    def __init__(self,data_path,transform=None):
        super(CustomDataset, self).__init__()
        self.data_path = data_path
        if transform is None:
            self.transform = transforms.Compose(
                [transforms.ToTensor(),  # 转化为tensor
                 transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))] # 归一化
            )
        else:
            self.transform = transform
        self.path_list = os.listdir(self.data_path)
    def __len__(self):
        return len(self.path_list)

    def __getitem__(self, idx: int):
        img_path = self.path_list[idx]
        label = int(img_path.split('.')[0].split('_')[0])
        label = torch.tensor(label, dtype=torch.int64)  # 把标签转换成int64
        img_path = os.path.join(self.data_path, img_path)  # 合成图片路径
        img = Image.open(img_path).convert('RGB')  # 读取图片,转化为RGB
        img = self.transform(img)  # 把图片转换成tensor
        label=torch.tensor(label, dtype=torch.int64)  # 把标签转换成int64
        return img, label
